# -*- coding: utf-8 -*-

__author__="Boris Avdeev"
__date__ ="$Mar 8, 2011 10:09:48 AM$"

import pymc as pm
import numpy as np
import os


class AgeModel(object):
    """Generic single grain and sample age model"""
    def __init__(self, tc, sample_age, RelErr=None, Age=None):
        """If Age or RelErr are None, they are modelled. Alternatively, constants or extermal models can be passed."""
        max_age=4500.
        name=tc.name
        ages = tc.quick_ages
        if RelErr is None:
            if sample_age: self.RelErr = pm.Uniform("RelErr_%s"%name, 0.05, 0.5, value=0.1)
            else: self.RelErr = 0.15  # Do not model error for single grain ages: no such information! 15% is the Caucasus average.
        else: self.RelErr = RelErr

        if Age is None:
            if sample_age: # assume common age for all grains
                init_age = np.mean(ages)
            else: # each grain has individual age
                init_age = ages
            self.Age = pm.Uniform("Age_%s"  % name, 0., max_age, value = init_age)
        else: self.Age = Age

        #TODO: DATA
        self.ObsAge = pm.Normal("ObsAge_%s" % name,  mu = self.Age, tau = 1./(self.RelErr*self.Age)**2, 
            value = ages, observed=True)

        # SIMULATION
        self.SimAge = pm.Lambda("SimAge_%s"%name,lambda mu=self.Age,re=self.RelErr: pm.rnormal(mu,1./(re*mu)**2,
            size=len(ages)),plot=False)


    def get_list(self):
        return self.__dict__.values()











class TCdata(object):
    """Base class for thermochron data"""
    def __init__(self):
        self.__sa_mcmc=None
        self.__sga_mcmc=None
        self.sa_model, self.sga_model  = None, None

    def sa_mcmc(self, update_length = 0):
        if self.sa_model==None: return None
        if self.__sa_mcmc==None:
            name = '.' + self.name + "_%s_SA"%self.method
            if not os.path.exists(name+'.hdf5'): update_length=5000
            db = pm.database.hdf5.load(name+'.hdf5')
            #update_length=5000
            self.__sa_mcmc = pm.MCMC(self.sa_model, db=db, name=name)
            self.__sa_mcmc.restore_sampler_state()
        if update_length>0:
            self.__sa_mcmc.sample(6*update_length/5, burn = update_length/5, thin = 0.001 * update_length)
            self.__sa_mcmc.save_state()
            self.__sa_mcmc.db.commit()
        return self.__sa_mcmc

    def sga_mcmc(self, update_length = 0):
        if self.sga_model==None: return None
        if self.__sga_mcmc==None:
            name = '.' + self.name + "_%s_SGA"%self.method
            if not os.path.exists(name+'.hdf5'): update_length=5000
            db = pm.database.hdf5.load(name + '.hdf5')
            #update_length = 5000                       #TODO: Remove this line once a bug(?) #351 in PyMC fixed.
            self.__sga_mcmc = pm.MCMC(self.sga_model, db = db, name=name)
            self.__sga_mcmc.restore_sampler_state()
        if update_length>0:
            print "Sampling single grain age model..."
            self.__sga_mcmc.sample(6*update_length/5, burn = update_length/5, thin = 0.001 * update_length)
            self.__sga_mcmc.save_state()
            self.__sga_mcmc.db.commit()
        return self.__sga_mcmc


    @property
    def sa_trace(self):
        if self.sa_model == None: return [None]
        return self.sa_mcmc().trace(self.sa_model.Age)[:]

    @property
    def sga_trace(self):
        if self.sga_model == None: return [None]
        return self.sga_mcmc().trace(self.sga_model.Age)[:]



    @property
    def sa(self):
        """Returns sample age and standard deviation estimates assuming all grains
        record the same age (i.e. record the same cooling history and have identical kinetics)"""
        return np.mean(self.sa_trace), np.std(self.sa_trace)

    @property
    def sga(self):
        """Single grain age means and standard deviations"""
        if self.sga_model == None: return (None, None)
        ages=[]
        traces = self.sga_trace
        for i in xrange(self.n_cnt):
            ages.append((np.mean(traces[:,i]), np.std(traces[:,i])))
        return ages


    @property
    def sa_hpd(self):
        return self.sa_mcmc(1000).Age.stats()['95% HPD interval'] #TODO remove 1000

    @property
    def sga_hpd(self):
        return self.sga_mcmc(1000).Age.stats()['95% HPD interval'] # TODO: Remove 5000

    @property
    def sga_mean(self):
        """Mean single grain ages"""
        if self.sga_model == None: return None
        ages=np.zeros(self.n_cnt)
        traces = self.sga_trace
        for i in xrange(self.n_cnt):
            ages[i] = np.mean(traces[:,i])
        return ages


